#include "mem_scan.h"
#include <TlHelp32.h>

size_t string_to_bytes(IN std::string sign, IN OUT short *dest) {
  sign.erase(std::remove(sign.begin(), sign.end(), ' '), sign.end());
  if (sign.size() % 2 != 0) {
    return 0;
  }
  size_t len = sign.size() / 2;
  ZeroMemory(dest, len);
  char *cstr_sign = const_cast<char *>(sign.c_str());
  for (size_t i = 0; i < len; i++) {
    // 先判断是不是 ?? 如果是, 映射为 0xFFFF
    if (cstr_sign[i * 2] == '?' && cstr_sign[i * 2 + 1] == '?') {
      dest[i] = (short)0xFFFF;
    } else {
      // 每两个字符就转为一个字节
      char c[] = {cstr_sign[i * 2], cstr_sign[i * 2 + 1], 0};
      dest[i] = (short)strtol(c, NULL, 16);
    }
  }

  return len;
}

bool cmp_sign(short *sign, byte *buf, size_t len) {
  for (size_t i = 0; i < len; i++) {
    if (sign[i] != (short)0xFFFF && sign[i] != (short)buf[i]) {
      return false;
    }
  }
  return true;
}

std::optional<std::pair<uint64, uint64>> get_module_addrs(uint32 pid) {
  // 创建模块信息相关结构变量
  MODULEENTRY32 mod_entry;
  mod_entry.dwSize = sizeof(MODULEENTRY32);
  // 创建进程快照
  HANDLE handle = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid);
  // 如果创建失败, 直接返回
  if (handle == INVALID_HANDLE_VALUE) {
    return std::nullopt;
  }
  // 对快照进行遍历
  bool ret = Module32First(handle, &mod_entry);
  while (ret) {
    // 根据名称进行比对
    // if (wcscmp(mod_name, mod_entry.szModule) == 0) {
    if (mod_entry.th32ProcessID == pid) {
      uint64 start_addr = (uint64)mod_entry.modBaseAddr;
      uint64 end_addr = start_addr + mod_entry.modBaseSize;
      CloseHandle(handle);
      return std::make_optional(std::make_pair(start_addr, end_addr));
    }
    ret = Module32Next(handle, &mod_entry);
  }
  CloseHandle(handle);
  return std::nullopt;
}

std::optional<uint64> get_module_base(uint32 pid) {
  auto addrs = get_module_addrs(pid);
  if (addrs.has_value()) {
    return addrs.value().first;
  } else {
    return std::nullopt;
  }
}

uint64 scan_raw_offset(HANDLE process, uint64 start, uint64 end,
                       std::string sign, int line_offset) {
  const uint32 page_size = 4096;
  byte page[page_size] = {0};
  size_t tmp = 0;
  byte cmd_bytes[0x10] = {0}; // 一般都是 4 字节的, 基址/call/大偏移都是4字节
  int cmd =
      0; // 指令数值, 基址/call/偏移等在机器码中的值, 可能是负数, 所以得用 int
  uint64 addr = 0; // 指令所在行的内存地址
  // 将 sign 转为 字节数组
  short sign_short[0x100] = {0};
  size_t sign_len = string_to_bytes(sign, sign_short);
  if (sign_len == 0) {
    return 0;
  }

  while (start < end) {
    bool success =
        ::ReadProcessMemory(process, (LPCVOID)start, page, page_size, NULL);
    if (!success)
      return 0;
    for (size_t i = 0; i < page_size; i++) {
      if (cmp_sign(sign_short, page + i, sign_len)) {
        addr = start + i + line_offset;
        return addr;
      }
      tmp = i;
    }
    start = start + tmp - (sign.size() / 2);
  }
  return uint64(0);
}

uint64 scan_raw_addr(HANDLE process, uint64 start, uint64 end, std::string sign,
                     int line_offset) {
  const uint32 page_size = 4096;
  byte page[page_size] = {0};
  size_t tmp = 0;
  byte cmdBytes[0x10] = {0}; // 一般都是 4 字节的, 基址/call/大偏移都是4字节
  int cmd =
      0; // 指令数值, 基址/call/偏移等在机器码中的值, 可能是负数, 所以得用 int
  uint64 addr = 0; // 指令所在行的内存地址
  // 将 sign 转为 字节数组
  short sign_short[0x100] = {0};
  size_t signLen = string_to_bytes(sign, sign_short);
  if (signLen == 0) {
    return 0;
  }

  while (start < end) {
    bool success =
        ::ReadProcessMemory(process, (LPCVOID)start, page, page_size, NULL);
    if (!success)
      return 0;
    for (size_t i = 0; i < page_size; i++) {
      if (cmp_sign(sign_short, page + i, signLen)) {
        addr = start + i + line_offset;
        return addr;
      }
      tmp = i;
    }
    start = start + tmp - (sign.size() / 2);
  }
  return uint64(0);
}

uint64 scan_raw_addr_rva(HANDLE process, uint64 start, uint64 end,
                         std::string sign, int line_offset) {
  auto addr = scan_raw_addr(process, start, end, sign, line_offset);
  if (addr > start) {
    return addr - start;
  }
  return 0;
}

/// @brief 通用定位函数, 返回定位到的 call/基址/偏移
/// @param process 游戏进程句柄
/// @param startAddr 游戏模块起始地址
/// @param endAddr 游戏模块结束地址
/// @param sign 特征码字符串 std::string 类型
/// @param lineOffset 行偏移, int 类型, 可以是负数
/// @param inlineOffset 行内偏移 只能是正数; call 一般是 1, 基址一般是
/// 4(有其他情况), 偏移就不确定了
/// @param cmdSize 定位数据指令字节数, 比如 48:8B0D 76D27502 inlineOffset 就是
/// 3, cmdSize 就是 76D27502 的字节数 4
/// @return 返回定位到的 call/基址/偏移
uint64 scan_raw(HANDLE process, uint64 start, uint64 end, std::string sign,
                int line_offset, uint32 inline_offset, uint32 cmd_len) {
  const uint32 page_size = 4096;
  byte page[page_size] = {0};
  size_t tmp = 0;
  byte cmd_bytes[0x10] = {0}; // 一般都是 4 字节的, 基址/call/大偏移都是4字节
  int cmd =
      0; // 指令数值, 基址/call/偏移等在机器码中的值, 可能是负数, 所以得用 int
  uint64 addr = 0; // 指令所在行的内存地址
  // 将 sign 转为 字节数组
  short sign_short[0x100] = {0};
  size_t sign_len = string_to_bytes(sign, sign_short);
  if (sign_len == 0) {
    return 0;
  }

  while (start < end) {
    bool success =
        ::ReadProcessMemory(process, (LPCVOID)start, page, page_size, NULL);
    if (!success)
      return 0;
    for (size_t i = 0; i < page_size; i++) {
      if (cmp_sign(sign_short, page + i, sign_len)) {
        addr = start + i + line_offset;
        // 从 指令所在行的内存地址 + 行内偏移 开始读, 读 cmdSize 字节到 cmdBytes
        // 数组中
        success = ::ReadProcessMemory(process, (LPCVOID)(addr + inline_offset),
                                      cmd_bytes, cmd_len, NULL);
        if (!success)
          return 0;
        memcpy(&cmd, cmd_bytes, cmd_len);
        // printf("addr: %llX; cmd: %X; inlineOffset %d; cmdSize: %d;\n", addr,
        // cmd, inlineOffset, cmdSize);
        return addr + cmd + inline_offset + cmd_len;
      }
      tmp = i;
    }
    start = start + tmp - (sign.size() / 2);
  }
  return uint64(0);
}

/// @brief 扫描 call. 调用 ScanRaw, inlineOffset 固定是 1, cmdSize 固定是 4,
/// 返回 rva
/// @param process
/// @param startAddr
/// @param endAddr
/// @param sign
/// @param lineOffset
/// @return rva
uint64 scan_call(HANDLE process, uint64 start, uint64 end, std::string sign,
                 int line_offset) {
  return scan_raw(process, start, end, sign, line_offset, 1, 4) - start;
}

/// @brief 扫描基址, 调用 ScanRaw, cmdSize 是固定 4 的, 返回 rva
/// @param process
/// @param startAddr
/// @param endAddr
/// @param sign
/// @param lineOffset
/// @param inlineOffset
/// @return rva
uint64 scan_base(HANDLE process, uint64 start, uint64 end, std::string sign,
                 int line_offset, uint32 inline_offset) {
  return scan_raw(process, start, end, sign, line_offset, inline_offset, 4) -
         start;
}

/// @brief 逻辑与 ScanRaw 一致, 但是偏移只需要直接返回指令的值即可.
/// @param process
/// @param startAddr
/// @param endAddr
/// @param sign
/// @param lineOffset
/// @param inlineOffset
/// @param cmdSize
/// @return
uint64 scan_offset(HANDLE process, uint64 start, uint64 end, std::string sign,
                   int line_offset, uint32 inline_offset, uint32 cmd_len) {
  const uint32 page_size = 4096;
  byte page[page_size] = {0};
  size_t tmp = 0;
  byte cmd_bytes[0x10] = {0}; // 一般都是 4 字节的, 基址/call/大偏移都是4字节
  uint64 cmd =
      0; // 指令数值, 基址/call/偏移等在机器码中的值, 这里可能是负数, 用 uint64
  uint64 addr = 0; // 指令所在行的内存地址
  // 将 sign 转为 字节数组
  short sign_short[0x100] = {0};
  size_t sign_len = string_to_bytes(sign, sign_short);
  if (sign_len == 0) {
    return 0;
  }

  while (start < end) {
    bool success =
        ::ReadProcessMemory(process, (LPCVOID)start, page, page_size, NULL);
    if (!success)
      return 0;
    for (size_t i = 0; i < page_size; i++) {
      if (cmp_sign(sign_short, page + i, sign_len)) {
        addr = start + i + line_offset;
        // 从 指令所在行的内存地址 + 行内偏移 开始读, 读 cmdSize 字节到 cmdBytes
        // 数组中
        success = ::ReadProcessMemory(process, (LPCVOID)(addr + inline_offset),
                                      cmd_bytes, cmd_len, NULL);
        if (!success)
          return 0;
        memcpy(&cmd, cmd_bytes, cmd_len);
        return cmd;
      }
      tmp = i;
    }
    start = start + tmp - (sign.size() / 2);
  }
  return uint64(0);
}

/// @brief 比较特殊的情况, 基址不好定位, 但是所在的 call 比较好定位, 通过 call
/// 的特征码定位到 call 再定位到基址.
/// @param process
/// @param startAddr
/// @param endAddr
/// @param sign
/// @param lineOffset
/// @param callLineOffset
/// @param callInlineOffset
/// @return rva
uint64 scan_base_in_call(HANDLE process, uint64 start, uint64 end,
                         std::string sign, int line_offset,
                         uint32 call_line_offset, uint32 call_inline_offset) {
  uint64 rva_start = start; // 保存模块基址, 用于计算最后结果的 rva
  const uint32 page_size = 4096;
  byte page[page_size] = {0};
  size_t tmp = 0;
  byte cmd_bytes[0x10] = {0}; // 一般都是 4 字节的, 基址/call/大偏移都是4字节
  int cmd = 0;     // 指令数值, 基址/call/偏移等在机器码中的值
  uint64 addr = 0; // 指令所在行的内存地址
  // 将 sign 转为 字节数组
  short shortSign[0x100] = {0};
  size_t sign_len = string_to_bytes(sign, shortSign);
  if (sign_len == 0) {
    return 0;
  }

  while (start < end) {
    bool success =
        ::ReadProcessMemory(process, (LPCVOID)start, page, page_size, NULL);
    if (!success)
      return 0;
    for (size_t i = 0; i < page_size; i++) {
      if (cmp_sign(shortSign, page + i, sign_len)) {
        addr = start + i + line_offset;
        // 从 指令所在行的内存地址 + 行内偏移 开始读, 读 cmdSize 字节到 cmdBytes
        // 数组中
        success = ::ReadProcessMemory(process, (LPCVOID)(addr + 1), cmd_bytes,
                                      4, NULL);
        if (!success)
          return 0;
        memcpy(&cmd, cmd_bytes, 4);
        uint64 callAddr = addr + cmd + 1 + 4;

        ZeroMemory(cmd_bytes, 0x10);
        cmd = 0;

        // call 内地址
        addr = callAddr + call_line_offset;
        // 从 指令所在行的内存地址 + 行内偏移 开始读, 读 4 字节到 cmdBytes
        // 数组中(基址都是4字节)
        success = ::ReadProcessMemory(
            process, (LPCVOID)(addr + call_inline_offset), cmd_bytes, 4, NULL);
        if (!success)
          return 0;
        memcpy(&cmd, cmd_bytes, 4);
        return addr + cmd + call_inline_offset + 4 - rva_start;
      }
      tmp = i;
    }
    start = start + tmp - (sign.size() / 2);
  }
  return uint64(0);
}
